import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping
from datetime import timedelta


# 1. 数据加载与预处理
def load_and_preprocess_data(file_path):
    # 读取数据
    df = pd.read_excel(file_path)

    # 转换日期格式
    df['date_id'] = pd.to_datetime(df['date_id'])
    df['new_date'] = pd.to_datetime(df['new_date'])

    # 计算生命周期特征
    df['days_since_new'] = (df['date_id'] - df['new_date']).dt.days

    # 按日期排序
    df.sort_values('date_id', inplace=True)

    # 创建时间序列特征
    df['day_of_week'] = df['date_id'].dt.dayofweek
    df['day_of_month'] = df['date_id'].dt.day
    df['month'] = df['date_id'].dt.month

    return df


# 2. 数据准备函数
def prepare_data(data, look_back=30):
    # 选择特征
    features = ['daily_sale_qty', 'days_since_new', 'day_of_week', 'day_of_month', 'month']
    dataset = data[features].values

    # 数据标准化
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(dataset)

    # 创建时间序列数据集
    X, y = [], []
    for i in range(look_back, len(scaled_data)):
        X.append(scaled_data[i - look_back:i, :])
        y.append(scaled_data[i, 0])  # 只预测销量

    X, y = np.array(X), np.array(y)

    return X, y, scaler


# 3. 构建LSTM模型
def build_lstm_model(input_shape):
    model = Sequential()
    model.add(LSTM(units=64, return_sequences=True, input_shape=input_shape))
    model.add(Dropout(0.2))
    model.add(LSTM(units=32, return_sequences=False))
    model.add(Dropout(0.2))
    model.add(Dense(units=1))
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model


# 4. 预测未来14天 - 修复版本
def forecast_future(model, last_sequence, scaler, future_days=14):
    predictions = []
    current_sequence = last_sequence.copy()

    # 生成未来日期的特征
    last_date = data['date_id'].iloc[-1]
    future_dates = [last_date + timedelta(days=i) for i in range(1, future_days + 1)]

    # 获取特征数量
    n_features = last_sequence.shape[1]

    for i in range(future_days):
        # 预测下一天
        pred = model.predict(current_sequence[np.newaxis, :, :])

        # 创建新数据点
        next_day = current_sequence[-1].copy()
        next_day[0] = pred[0, 0]  # 更新销量

        # 更新日期特征
        next_date = future_dates[i]
        next_day[2] = next_date.dayofweek / 6.0  # 标准化星期
        next_day[3] = next_date.day / 31.0  # 标准化日期
        next_day[4] = next_date.month / 12.0  # 标准化月份

        # 更新生命周期
        # 先将当前的生命周期逆变换回原始值
        current_life = next_day[1] * (scaler.data_max_[1] - scaler.data_min_[1]) + scaler.data_min_[1]
        # 加一天
        current_life += 1
        # 再标准化
        next_day[1] = (current_life - scaler.data_min_[1]) / (scaler.data_max_[1] - scaler.data_min_[1])

        # 更新序列
        current_sequence = np.vstack([current_sequence[1:], next_day])

        # 保存预测结果
        predictions.append(pred[0, 0])

    # 修复逆标准化问题：创建与原始特征数量相同的矩阵
    # 创建全零矩阵，维度为 (预测天数, 特征数量)
    dummy = np.zeros((len(predictions), n_features))
    # 将预测结果放入第一列（销量列）
    dummy[:, 0] = np.array(predictions).reshape(-1)
    # 逆标准化整个矩阵
    predictions = scaler.inverse_transform(dummy)[:, 0]

    return future_dates, predictions


# 主程序
if __name__ == "__main__":
    # 1. 加载数据
    file_path = 'sales.xlsx'
    data = load_and_preprocess_data(file_path)

    # 2. 准备训练数据
    look_back = 60  # 使用过去60天预测未来
    X, y, scaler = prepare_data(data, look_back)

    # 3. 划分训练集
    train_size = int(len(X) * 0.8)
    X_train, X_val = X[:train_size], X[train_size:]
    y_train, y_val = y[:train_size], y[train_size:]

    # 4. 构建并训练模型
    model = build_lstm_model((X_train.shape[1], X_train.shape[2]))
    early_stop = EarlyStopping(monitor='val_loss', patience=10)
    history = model.fit(
        X_train, y_train,
        epochs=5,
        batch_size=32,
        validation_data=(X_val, y_val),
        callbacks=[early_stop],
        verbose=1
    )

    # 5. 预测未来14天
    last_sequence = X[-1]  # 获取最后60天的数据
    future_dates, predictions = forecast_future(model, last_sequence, scaler)

    # 6. 输出结果
    total_sales = sum(predictions)
    print("\n未来14天销量预测：")
    for date, qty in zip(future_dates, predictions):
        print(f"{date.strftime('%Y-%m-%d')}: {qty:.2f}件")

    print(f"\n未来14天总销量预测: {total_sales:.2f}件")

    # 7. 可视化结果
    # plt.figure(figsize=(12, 6))
    # plt.plot(data['date_id'], data['daily_sale_qty'], label='historical sales')
    # plt.plot(future_dates, predictions, 'ro-', label='predict sales')
    # plt.title('销量预测结果')
    # plt.xlabel('日期')
    # plt.ylabel('销量')
    # plt.legend()
    # plt.grid(True)
    # plt.show()